from threading import Condition, RLock
import numpy as np


class Stack(object):
    def __init__(self, max_size=-1):
        self._li = []
        self._max_size = max_size

    def empty(self):
        return self.size() == 0

    def push(self, o):
        if 0 < self._max_size <= self.size():
            raise RuntimeError('The stack is full.')
        self._li.append(o)

    def pop(self):
        return self._li.pop(-1)

    def peek(self):
        return self._li[-1]

    def clear(self):
        self._li.clear()

    def size(self):
        return len(self._li)


class Queue(object):
    def __init__(self, max_size=-1):
        self._li = []
        self._max_size = max_size

    def enqueue(self, o):
        if 0 < self._max_size <= self.size():
            raise RuntimeError('The queue is full.')
        self._li.append(o)

    def dequeue(self):
        return self._li.pop(0)

    def empty(self):
        return self.size() == 0

    def peek(self):
        return self._li[0]

    def clear(self):
        self._li.clear()

    def size(self):
        return len(self._li)


class CountDownLatch:
    def __init__(self, count: int):
        if count <= 0:
            raise ValueError('count should larger than 0.')
        self._count = count
        self._cond = Condition()

    def wait(self):
        with self._cond:
            while self._count > 0:
                self._cond.wait()

    def count_down(self):
        with self._cond:
            self._count -= 1
            if self._count == 0:
                self._cond.notify_all()


def cos_similarity(vec1, vec2, dtype=np.float64):
    vec1 = np.array(vec1, dtype=dtype)
    vec2 = np.array(vec2, dtype=dtype)

    return vec1.dot(vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))


def pearson_similarity(vec1: list, vec2: list, dtype=np.float64):
    vec1 = np.array(vec1, dtype=dtype)
    vec2 = np.array(vec2, dtype=dtype)
    vec1 -= vec1.mean(dtype=np.float64)
    vec1 -= vec2.mean(dtype=dtype)
    return cos_similarity(vec1, vec2)


class SyncDict(object):
    def __init__(self):
        self._dict = {}
        self._lock = RLock()

    def _contains(self, key):
        return key in self._dict.keys()

    def contains(self, key):
        with self._lock:
            return self._contains(key)

    def _get(self, key):
        return self._dict[key]

    def get_or_default(self, key, default):
        with self._lock:
            if self._contains(key):
                return self._get(key)
            return default

    def get(self, key):
        with self._lock:
            return self._get(key)

    def _put(self, key, val):
        self._dict[key] = val

    def put(self, key, val):
        with self._lock:
            self._put(key, val)

    def put_just_not_exists(self, key, default):
        with self._lock:
            if not self._contains(key):
                self._put(key, default)

    def _size(self):
        return len(self._dict)

    def size(self):
        with self._lock:
            return self._size()

    def is_empty(self):
        with self._lock:
            return self._size() == 0

    def clear(self):
        with self._lock:
            self._dict.clear()

    def _remove(self, key):
        if self._contains(key):
            return self._dict.pop(key)
        return None

    def remove(self, key):
        with self._lock:
            return self._remove(key)

    def modify_if(self, cond_func, new_val):
        """
        cond_func是形如function(key, val)的函数
        会遍历dict中的所有kv，调用cond_func，当返回真的时候就会修改这个key对应的val为new_val
        """
        with self._lock:
            for key in self._dict.keys():
                val = self._get(key)
                if cond_func(key, val):
                    self._put(key, new_val)

    def remove_if(self, cond_func):
        """
        cond_func是形如function(key, val)的函数
        会遍历dict中的所有kv，调用cond_func，当返回真的时候就会移除这个key
        """
        with self._lock:
            keys = []
            for key in self._dict.keys():
                val = self._get(key)
                if cond_func(key, val):
                    keys.append(key)
            for key in keys:
                self._remove(key)

    def compute(self, key, cpt_func):
        with self._lock:
            if self._contains(key):
                val = self._get(key)
                new_val = cpt_func(key, val)
                self._put(key, new_val)
